package com.aabte.nettyleran.netty.http.websocket;

import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.codec.http.*;
import io.netty.handler.codec.http.websocketx.*;
import io.netty.handler.stream.ChunkedWriteHandler;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.util.Date;

import static io.netty.handler.codec.http.HttpHeaderUtil.isKeepAlive;

/**
 * @author Daniel
 * @version 1.0
 * @date 2020/3/14
 */
public class WebSocketServer {

    public static void main(String[] args) throws InterruptedException {
        int port = 8080;
        if (args != null && args.length > 0) {
            try {
                Integer.valueOf(args[0]);
            } catch (NumberFormatException e) {
                e.printStackTrace();
            }
        }

        new WebSocketServer().run(port);
    }

    private void run(int port) throws InterruptedException {
        EventLoopGroup bossGroup = new NioEventLoopGroup();
        EventLoopGroup workerGroup = new NioEventLoopGroup();

        ServerBootstrap serverBootstrap = new ServerBootstrap();
        serverBootstrap.group(bossGroup, workerGroup)
                .channel(NioServerSocketChannel.class)
                .childHandler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) throws Exception {
                        ch.pipeline().addLast("http-codec", new HttpServerCodec());
                        ch.pipeline().addLast("aggregator", new HttpObjectAggregator(65535));
                        ch.pipeline().addLast("http-chunked", new ChunkedWriteHandler());
                        ch.pipeline().addLast("handler", new WebSocketServerHandler());
                    }
                });

        try {
            Channel channel = serverBootstrap.bind(port).sync().channel();
            System.out.println("Web socket server started at port: " + port);
            System.out.println("Open your browser and navigate to http://localhost:" + port + "/");

            channel.closeFuture().sync();
        } finally {
            bossGroup.shutdownGracefully();
            workerGroup.shutdownGracefully();
        }
    }

    @Slf4j
    private static class WebSocketServerHandler extends SimpleChannelInboundHandler<Object> {

        private WebSocketServerHandshaker handshaker;

        @Override
        protected void messageReceived(ChannelHandlerContext ctx, Object msg) throws Exception {
            if (msg instanceof FullHttpRequest) {
                // 传统HTTP接入
                handleHttpRequest(ctx, (FullHttpRequest) msg);
            } else if (msg instanceof WebSocketFrame) {
                // WebSocket接入
                handleWebSocketFrame(ctx, (WebSocketFrame) msg);
            }
        }

        @Override
        public void channelReadComplete(ChannelHandlerContext ctx) throws Exception {
            ctx.flush();
        }

        private void handleHttpRequest(ChannelHandlerContext ctx, FullHttpRequest httpRequest) {
            if (!httpRequest.decoderResult().isSuccess()
                    || (!"websocket".equalsIgnoreCase(httpRequest.headers().getAndConvert("Upgrade")))) {
                // HTTP解码失败，返回HTTP异常
                sendHttpResponse(ctx, httpRequest, new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.BAD_REQUEST));
                return;
            }

            // 构造握手响应返回，本机测试
            WebSocketServerHandshakerFactory handshakerFactory = new WebSocketServerHandshakerFactory("ws://localhost:8080/websocket", null, false);
            handshaker = handshakerFactory.newHandshaker(httpRequest);
            if (null == handshaker) {
                WebSocketServerHandshakerFactory.sendUnsupportedVersionResponse(ctx.channel());
            } else {
                handshaker.handshake(ctx.channel(), httpRequest);
            }
        }

        private void sendHttpResponse(ChannelHandlerContext ctx, FullHttpRequest fullHttpRequest, DefaultFullHttpResponse fullHttpResponse) {
            // 返回应答给客户端
            if (fullHttpResponse.status().code() != 200) {
                ByteBuf byteBuf = Unpooled.copiedBuffer(fullHttpResponse.status().toString(), StandardCharsets.UTF_8);
                fullHttpResponse.content().writeBytes(byteBuf);
                byteBuf.release();
                HttpHeaderUtil.setContentLength(fullHttpResponse, fullHttpResponse.content().readableBytes());
            }

            ChannelFuture channelFuture = ctx.channel().writeAndFlush(fullHttpResponse);
            if (!isKeepAlive(fullHttpRequest) || fullHttpResponse.status().code() != 200) {
                // 非Keep-Alive，关闭连接
                channelFuture.addListener(ChannelFutureListener.CLOSE);
            }
        }

        private void handleWebSocketFrame(ChannelHandlerContext ctx, WebSocketFrame webSocketFrame) {
            if (webSocketFrame instanceof CloseWebSocketFrame) {
                // 是关闭链路的指令
                handshaker.close(ctx.channel(), (CloseWebSocketFrame) webSocketFrame.retain());
                return;
            }

            if (webSocketFrame instanceof PingWebSocketFrame) {
                // 是Ping消息
                ctx.channel().write(new PongWebSocketFrame(webSocketFrame.content().retain()));
                return;
            }

            if (!(webSocketFrame instanceof TextWebSocketFrame)) {
                // 仅支持文本消息，不进行二进制消息
                throw new UnsupportedOperationException(String.format("%s frame types not supported", webSocketFrame.getClass().getName()));
            }

            // 返回应答消息
            String request = ((TextWebSocketFrame) webSocketFrame).text();
            log.info("{} received {}", ctx.channel(), request);
            ctx.channel().write(new TextWebSocketFrame(request + ", 欢迎使用Netty WebSocket服务，现在时间: " + new Date().toString()));
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            cause.printStackTrace();
            ctx.close();
        }
    }
}
